import sqlite3
import torch
import pickle
import threading
from bert_finetune.utils import DualRelationProcessor

class EmbeddingSearchEngine:
    def __init__(self, model_dir, db_path, max_length=128):
        self.processor = DualRelationProcessor.from_pretrained(model_dir, max_length)
        self.db_path = db_path
        self.local = threading.local()  # 线程局部存储
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def _get_db_connection(self):
        if not hasattr(self.local, "db_connection"):
            self.local.db_connection = sqlite3.connect(
                self.db_path,
                check_same_thread=False
            )
        return self.local.db_connection

    def close(self):
        if hasattr(self.local, "db_connection"):
            self.local.db_connection.close()
            delattr(self.local, "db_connection")

    def get_query_embedding(self, query):
        emb = self.processor.get_hobby_embedding(query)
        if not isinstance(emb, torch.Tensor):
            emb = torch.tensor(emb, dtype=torch.float32)
        return emb.to(self.device)

    def search(self, query, top_k=5):
        query_emb = self.get_query_embedding(query)
        conn = self._get_db_connection()
        cursor = conn.cursor()
        
        try:
            cursor.execute("""
                SELECT id, full_name, embedding 
                FROM scenic_spots 
                WHERE embedding IS NOT NULL
            """)
            attractions = cursor.fetchall()

            if not attractions:
                return []

            results = []
            for attr in attractions:
                attr_id, name, emb_blob = attr
                try:
                    # 反序列化并转换为Tensor
                    emb_tensor = pickle.loads(emb_blob)
                    if not isinstance(emb_tensor, torch.Tensor):
                        emb_tensor = torch.tensor(emb_tensor, dtype=torch.float32)
                    emb_tensor = emb_tensor.to(self.device)

                    # 计算相似度分数
                    score = self.processor.get_hobby_attraction_score_from_embeddings(
                        query_emb, emb_tensor
                    )

                    # 关键修复：如果是Tensor才调用item()，否则直接使用（兼容float类型）
                    if isinstance(score, torch.Tensor):
                        similarity_score = score.item()  # Tensor转float
                    else:
                        similarity_score = score  # 已经是float，直接使用

                    results.append({
                        "attraction_id": attr_id,
                        "name": name,
                        "similarity_score": similarity_score
                    })
                except Exception as e:
                    print(f"处理景点 {attr_id} 出错: {e}")
                    continue

            results.sort(key=lambda x: x["similarity_score"], reverse=True)
            return results[:top_k]

        except Exception as e:
            print(f"搜索逻辑出错: {e}")
            return []
        finally:
            cursor.close()

    def search_by_embedding(self, input_embedding, top_k=10):
        """
        根据输入嵌入嵌入向量从数据库检索余弦相似性最高的前10个景点
        """
        # 确保输入嵌入在正确设备上
        input_emb = input_embedding.to(self.device)
        conn = self._get_db_connection()
        cursor = conn.cursor()
        
        try:
            cursor.execute("""
                SELECT id, full_name, embedding 
                FROM scenic_spots 
                WHERE embedding IS NOT NULL
            """)
            attractions = cursor.fetchall()

            if not attractions:
                return []

            results = []
            for attr in attractions:
                attr_id, name, emb_blob = attr
                try:
                    # 反序列化并转换为Tensor
                    emb_tensor = pickle.loads(emb_blob)
                    if not isinstance(emb_tensor, torch.Tensor):
                        emb_tensor = torch.tensor(emb_tensor, dtype=torch.float32)
                    emb_tensor = emb_tensor.to(self.device)

                    # 计算余弦相似度分数
                    similarity_score = self.processor.get_cosine_similarity(input_emb, emb_tensor)

                    results.append({
                        "attraction_id": attr_id,
                        "name": name,
                        "similarity_score": similarity_score
                    })
                except Exception as e:
                    print(f"处理景点 {attr_id} 出错: {e}")
                    continue

            results.sort(key=lambda x: x["similarity_score"], reverse=True)
            return results[:top_k]

        except Exception as e:
            print(f"嵌入搜索逻辑出错: {e}")
            return []
        finally:
            cursor.close()
    